# 实现接触相关的特征提取，例如接触区域
import numpy as np  
def get_contact_region(position: np.ndarray, forces: np.ndarray, method: str = "force_thres", **kwargs) -> np.ndarray:
    if method == "force_thres":
        threshold = kwargs.get("threshold")
        return _get_contact_region_by_threshold(forces, threshold)
    elif method == "soft_thres":
        threshold = kwargs.get("threshold")
        scale = kwargs.get("scale", 600.0)
        return _get_contact_region_by_soft_threshold(forces, threshold, scale)
    else:
        raise ValueError("Unknown method: {}".format(method))


def _get_contact_region_by_threshold(forces: np.ndarray, threshold: float) -> np.ndarray:
    '''
    Get the contact region based on force thresholding.
    Args:
        forces (np.ndarray): 3D Forces of shape (N, 400, 3) or (N, 20,20,3).
        threshold (float): The force threshold.
    Returns:
        np.ndarray: A boolean mask indicating the contact region.
    '''
    force_norm = np.linalg.norm(forces, axis=-1) # (N, 400) or (N, 20,20)
    contact_mask = force_norm > threshold # (N, 400) or (N, 20,20)
    return contact_mask

def _get_contact_region_by_soft_threshold(forces: np.ndarray, threshold: float, scale: float = 600.0) -> np.ndarray:
    '''
    Get the contact region based on soft force thresholding.
    Args:
        forces (np.ndarray): 3D Forces of shape (N, 400, 3) or (N, 20,20,3).
        threshold (float): The force threshold.
        scale (float): The scaling factor for the sigmoid function.
    Returns:
        np.ndarray: A float mask indicating the contact region with values between 0 and 1.
    '''
    force_norm = np.linalg.norm(forces, axis=-1) # (N, 400) or (N, 20,20)
    contact_mask = 1 / (1 + np.exp(-scale * (force_norm - threshold))) # (N, 400) or (N, 20,20)
    return contact_mask